import os
import numpy as np
from plyfile import PlyData, PlyElement
import shutil


def standardize_bbox(pcl):
    mins = np.amin(pcl, axis=0)
    maxs = np.amax(pcl, axis=0)
    center = (mins + maxs) / 2.
    scale = np.amax(maxs - mins)
    #print("Center: {}, Scale: {}".format(center, scale))
    result = ((pcl - center) / scale).astype(np.float32)  # [-0.5, 0.5]
    return result


xml_head = \
    """
<scene version="0.5.0">
    <integrator type="path">
        <integer name="maxDepth" value="-1"/>
    </integrator>
    <sensor type="perspective">
        <float name="farClip" value="100"/>
        <float name="nearClip" value="0.1"/>
        <transform name="toWorld">
            <lookat origin="4,0,0" target="0,0,0" up="0,0,1"/>
        </transform>
        <float name="fov" value="25"/>
        
        <sampler type="ldsampler">
            <integer name="sampleCount" value="256"/>
        </sampler>
        <film type="ldrfilm">
            <integer name="width" value="1600"/>
            <integer name="height" value="1200"/>
            <rfilter type="gaussian"/>
            <boolean name="banner" value="false"/>
        </film>
    </sensor>
    
    <bsdf type="roughplastic" id="surfaceMaterial">
        <string name="distribution" value="ggx"/>
        <float name="alpha" value="0.05"/>
        <float name="intIOR" value="1.46"/>
        <rgb name="diffuseReflectance" value="1,1,1"/> <!-- default 0.5 -->
    </bsdf>
    
"""

xml_shape = \
    """
    <shape type="ply">
        <string name="filename" value="{}"/>
        <bsdf type="diffuse">
            <rgb name="reflectance" value="{},{},{}"/>
        </bsdf>
    </shape>
    """

xml_tail = \
    """
    <shape type="rectangle">
        <ref name="bsdf" id="surfaceMaterial"/>
        <transform name="toWorld">
            <scale x="10" y="10" z="1"/>
            <translate x="0" y="0" z="-0.5"/>
        </transform>
    </shape>
    
    <shape type="rectangle">
        <transform name="toWorld">
            <scale x="10" y="10" z="1"/>
            <rotate y="1" angle="90"/>
            <translate x="-5" y="0" z="0"/>
        </transform>
        <emitter type="area">
            <rgb name="radiance" value="6,6,6"/>
        </emitter>
    </shape>
</scene>
"""

file_path = '/home/ljl/hdd/lidarcap_result_gallery_video/'
xml_file_path = '/home/ljl/hdd/lidarcap_video/lidarcap_result_gallery_video_mesh_xml/'
dirs = [filename for filename in os.listdir(file_path)]
idxxx = 0
for dir in dirs:
    file_dirs = [filename for filename in os.listdir(file_path + dir + '/')]
    for file_dir in file_dirs:

        filenames = [filename for filename in os.listdir(file_path + dir + '/' + file_dir)]

        if os.path.exists(xml_file_path + dir + '/' + file_dir + '/' ):
            shutil.rmtree(xml_file_path + dir + '/' + file_dir + '/' )
        os.makedirs(xml_file_path + dir + '/' + file_dir + '/' )
        for filename in filenames:
            if filename.endswith('mesh.ply'):
                idxxx = idxxx + 1
                print(idxxx)
                color_chose = 0
                ## 0  blue : stgcn method
                ## 1  red
                ## 2  purple
                ## 3  yellow
                ## 4  green

                # file_index = '7260_attention_stgcn_mesh'
                file_index = filename[0:-4]
                in_filename = file_path + dir + '/' + file_dir + '/' + filename



                out_filename = os.path.join(os.path.dirname(xml_file_path + dir + '/' + file_dir + '/' ), file_index + '_stand.ply')

                xmlname = xml_file_path + dir + '/' + file_dir + '/' + file_index + '.xml'
                ply_data = PlyData.read(in_filename)
                vertices_data = ply_data['vertex'].data
                points = np.array([[x, y, z] for x, y, z in vertices_data])

                theta_x = 0
                theta_y = 0
                theta_z = 180 / 180 * np.pi

                Rot_x = np.identity(3)
                Rot_y = np.identity(3)
                Rot_z = np.identity(3)

                Rot_x[1][1] = np.cos(theta_x)
                Rot_x[1][2] = -np.sin(theta_x)
                Rot_x[2][1] = np.sin(theta_x)
                Rot_x[2][2] = np.cos(theta_x)

                Rot_y[0][0] = np.cos(theta_y)
                Rot_y[0][2] = np.sin(theta_y)
                Rot_y[2][0] = -np.sin(theta_y)
                Rot_y[2][2] = np.cos(theta_y)

                Rot_z[0][0] = np.cos(theta_z)
                Rot_z[0][1] = -np.sin(theta_z)
                Rot_z[1][0] = np.sin(theta_z)
                Rot_z[1][1] = np.cos(theta_z)

                points = np.matmul(np.matmul(np.matmul(points, Rot_x.T), Rot_y.T), Rot_z.T)
                points = standardize_bbox(points)
                points[:, 2] += 0.1
                points = [(points[i, 0], points[i, 1], points[i, 2])
                          for i in range(points.shape[0])]
                vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')])
                vertex = PlyElement.describe(vertex, 'vertex', comments=['vertices'])
                face = PlyElement.describe(ply_data['face'].data, 'face', comments=['faces'])
                PlyData([vertex, face]).write(out_filename)

                xml_segments = [xml_head]


                if color_chose == 0:
                    # Blue stgcn method color
                    xml_segments.append(xml_shape.format(
                        out_filename, 139 / 255, 161 / 255, 199 / 255))
                elif color_chose == 1:
                    ## red color
                    xml_segments.append(xml_shape.format(
                        out_filename, 161 / 255, 87 / 255, 78 / 255
                    ))
                elif color_chose == 2:
                    ## purple color
                    xml_segments.append(xml_shape.format(
                        out_filename, 141 / 255, 115 / 255, 164 / 255
                    ))
                elif color_chose == 3:
                    ## yellow color
                    xml_segments.append(xml_shape.format(
                      out_filename, 204/255, 183/255, 92 / 255
                    ))
                elif color_chose == 4:
                    # Green color
                    xml_segments.append(xml_shape.format(
                        out_filename, 147 / 255, 179 / 255, 107 / 255
                    ))

                xml_segments.append(xml_tail)
                xml_content = str.join('', xml_segments)

                with open(xmlname, 'w') as f:
                    f.write(xml_content)
